// gpucore/buffer.go
package gpucore

import (
	"fmt"

	"github.com/cogentcore/webgpu/wgpu"
)

// GPUBuffer 表示 GPU 上的一块内存
type GPUBuffer interface {
	Read() ([]float32, error)
	ReadAsUint32() ([]uint32, error) // 新增方法
	Raw() *wgpu.Buffer
	Release()
	Size() uint64
}

type gpuBuffer struct {
	device *wgpu.Device
	buffer *wgpu.Buffer
	size   uint64
}

func newGPUBuffer(device *wgpu.Device, buf *wgpu.Buffer, size uint64) GPUBuffer {
	return &gpuBuffer{device: device, buffer: buf, size: size}
}

func (b *gpuBuffer) Raw() *wgpu.Buffer { return b.buffer }
func (b *gpuBuffer) Size() uint64      { return b.size }
func (b *gpuBuffer) Release() {
	if b.buffer != nil {
		b.buffer.Release()
		b.buffer = nil
	}
}

func (b *gpuBuffer) Read() ([]float32, error) {
	if b.buffer == nil {
		return nil, fmt.Errorf("buffer is nil")
	}
	staging, err := b.device.CreateBuffer(&wgpu.BufferDescriptor{
		Size:  b.size,
		Usage: wgpu.BufferUsageMapRead | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("create staging buffer: %w", err)
	}
	defer staging.Release()
	encoder, err := b.device.CreateCommandEncoder(nil)
	if err != nil {
		return nil, fmt.Errorf("create command encoder: %w", err)
	}
	encoder.CopyBufferToBuffer(b.buffer, 0, staging, 0, b.size)
	cmd, err := encoder.Finish(nil)
	encoder.Release()
	if err != nil {
		return nil, fmt.Errorf("finish command encoder: %w", err)
	}
	defer cmd.Release()
	b.device.GetQueue().Submit(cmd)
	done := make(chan error, 1)
	if err := staging.MapAsync(wgpu.MapModeRead, 0, b.size, func(status wgpu.BufferMapAsyncStatus) {
		done <- mapStatusToError(status)
	}); err != nil {
		return nil, fmt.Errorf("MapAsync: %w", err)
	}
	b.device.Poll(true, nil)
	if err := <-done; err != nil {
		return nil, err
	}
	defer staging.Unmap()
	raw := staging.GetMappedRange(0, uint(b.size))
	return wgpu.FromBytes[float32](raw), nil
}

// 增加用于调试的原始字节读取
func (b *gpuBuffer) ReadAsDebug() ([]byte, error) {
	if b.buffer == nil {
		return nil, fmt.Errorf("buffer is nil")
	}
	staging, err := b.device.CreateBuffer(&wgpu.BufferDescriptor{
		Size:  b.size,
		Usage: wgpu.BufferUsageMapRead | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("create staging buffer: %w", err)
	}
	defer staging.Release()
	encoder, err := b.device.CreateCommandEncoder(nil)
	if err != nil {
		return nil, fmt.Errorf("create command encoder: %w", err)
	}
	encoder.CopyBufferToBuffer(b.buffer, 0, staging, 0, b.size)
	cmd, err := encoder.Finish(nil)
	encoder.Release()
	if err != nil {
		return nil, fmt.Errorf("finish command encoder: %w", err)
	}
	defer cmd.Release()
	b.device.GetQueue().Submit(cmd)
	done := make(chan error, 1)
	if err := staging.MapAsync(wgpu.MapModeRead, 0, b.size, func(status wgpu.BufferMapAsyncStatus) {
		done <- mapStatusToError(status)
	}); err != nil {
		return nil, fmt.Errorf("MapAsync: %w", err)
	}
	b.device.Poll(true, nil)
	if err := <-done; err != nil {
		return nil, err
	}
	defer staging.Unmap()
	return staging.GetMappedRange(0, uint(b.size)), nil
}

// ReadAsUint32 从 GPU buffer 读取为 []uint32（用于 pivot 等整数数据）
func (b *gpuBuffer) ReadAsUint32() ([]uint32, error) {
	if b.buffer == nil {
		return nil, fmt.Errorf("buffer is nil")
	}
	staging, err := b.device.CreateBuffer(&wgpu.BufferDescriptor{
		Size:  b.size,
		Usage: wgpu.BufferUsageMapRead | wgpu.BufferUsageCopyDst,
	})
	if err != nil {
		return nil, fmt.Errorf("create staging buffer: %w", err)
	}
	defer staging.Release()
	encoder, err := b.device.CreateCommandEncoder(nil)
	if err != nil {
		return nil, fmt.Errorf("create command encoder: %w", err)
	}
	encoder.CopyBufferToBuffer(b.buffer, 0, staging, 0, b.size)
	cmd, err := encoder.Finish(nil)
	encoder.Release()
	if err != nil {
		return nil, fmt.Errorf("finish command encoder: %w", err)
	}
	defer cmd.Release()
	b.device.GetQueue().Submit(cmd)
	done := make(chan error, 1)
	if err := staging.MapAsync(wgpu.MapModeRead, 0, b.size, func(status wgpu.BufferMapAsyncStatus) {
		done <- mapStatusToError(status)
	}); err != nil {
		return nil, fmt.Errorf("MapAsync: %w", err)
	}
	b.device.Poll(true, nil)
	if err := <-done; err != nil {
		return nil, err
	}
	defer staging.Unmap()
	raw := staging.GetMappedRange(0, uint(b.size))
	return wgpu.FromBytes[uint32](raw), nil
}
